
#include "replication_pad2d_tiling.h"
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"
#include <algorithm>

namespace optiling {
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{

  ReplicationPad2dTilingData tiling;

  auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo()); 
  auto aivNum = ascendcPlatform.GetCoreNum();

  auto shape_x = context->GetInputTensor(0)->GetOriginShape();
  auto shape_dim = context->GetInputTensor(1)->GetOriginShape();
  int32_t pad_shapeNum = shape_dim.GetDimNum();
  //std::cout<<"shape_x:"<<shape_x<<std::endl;
  int32_t pad_dim1;
  int32_t pad_dim2;
  int32_t pad_dim3;
  if(pad_shapeNum == 1){
    pad_dim1 = shape_dim.GetDim(0);

    //std::cout<<"pad_dim1:"<<pad_dim1<<std::endl;
  }
  else if(pad_shapeNum == 2){
    pad_dim1 = shape_dim.GetDim(0);
    std::cout<<"pad_dim1:"<<pad_dim1<<std::endl;
    pad_dim2 = shape_dim.GetDim(1);

  }
  else if(pad_shapeNum == 3){
    pad_dim1 = shape_dim.GetDim(0);
    pad_dim2 = shape_dim.GetDim(1);
    pad_dim3 = shape_dim.GetDim(2);
  }
  // pad_dim = shape_dim.GetDim(0);
//   auto shape_dim = context->GetInputTensor(1)->GetOriginShape();

  // if(shape_dim.GetDimNum() ==1){
    // int32_t dim1 = shape_dim.GetDim(0)/2; //左
    // std::cout<<"dim1:"<<dim1<<std::endl;
    // int32_t dim2 = shape_dim.GetDim(0)/2; //右
    // std::cout<<"dim2:"<<dim2<<std::endl;
    // int32_t dim3 = shape_dim.GetDim(0)/2; //上
    // std::cout<<"dim3:"<<dim3<<std::endl;
    // int32_t dim4 = shape_dim.GetDim(0)/2; //下
    // std::cout<<"dim4:"<<dim4<<std::endl;
  // }
  // else{
  //   int32_t dim1 = shape_dim.GetDim(0); //Pading的第一个维度  左
  //   std::cout<<"dim1:"<<dim1<<std::endl;
  //   int32_t dim2 = shape_dim.GetDim(1); //Pading的第二个维度  右
  //   std::cout<<"dim2:"<<dim2<<std::endl;
  //   int32_t dim3 = shape_dim.GetDim(2); //Pading的第三个维度  上
  //   std::cout<<"dim3:"<<dim3<<std::endl;
  //   int32_t dim4 = shape_dim.GetDim(3); //Pading的第四个维度  下
  //   std::cout<<"dim4:"<<dim4<<std::endl;
  // }

  

  int32_t M = 1, N = 1, Z = 1;
  if(shape_x.GetDimNum() > 2)
  {

    for (int i = 0; i < shape_x.GetDimNum() - 2; i++)

        Z *= shape_x.GetDim(i);
        M = shape_x.GetDim(shape_x.GetDimNum() - 2);
        N = shape_x.GetDim(shape_x.GetDimNum() - 1);

        std::cout<<"M:"<<M<<std::endl;
        std::cout<<"N:"<<N<<std::endl;
        std::cout<<"Z:"<<Z<<std::endl;
        // Z = shape_x.GetDim(0);
    }
  else
    {
        M = shape_x.GetDim(0);
        N = shape_x.GetDim(1);
        Z = 1;
    }

  // int32_t new_y_nums = Z *(M+dim1+dim2)*(N+dim3+dim4); 
  tiling.set_M(M);
  tiling.set_N(N);
  tiling.set_Z(Z);
  
  tiling.set_pad_dim1(pad_dim1);
  tiling.set_pad_dim2(pad_dim2);
  tiling.set_pad_dim3(pad_dim3);
  tiling.set_pad_shapeNum(pad_shapeNum);
//   tiling.set_dim1(dim1);
//   tiling.set_dim2(dim2);
//   tiling.set_dim3(dim3);
//   tiling.set_dim4(dim4);

  // std::cout<<"new_y_nums:"<<new_y_nums<<std::endl;

  // tiling.set_new_y_nums(new_y_nums); 

  context->SetBlockDim(1);
  tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
  context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());

  return ge::GRAPH_SUCCESS;
}
}


namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}
}


namespace ops {
class ReplicationPad2d : public OpDef {
public:
    explicit ReplicationPad2d(const char* name) : OpDef(name)
    {
        this->Input("x")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("paddings")
            .ParamType(REQUIRED)
            .DataType({ge::DT_INT32, ge::DT_INT32})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});

        this->SetInferShape(ge::InferShape);

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend310b");
        //this->AICore().AddConfig("ascend310p");

    }
};

OP_ADD(ReplicationPad2d);
}
